/*
 * Copyright 2024 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "nlphelper.h"

#include <iostream>
#include <future>
#include <filesystem>

#include <kylin-ai/ai-base/private/osassistant.h>
#include <jsoncpp/json/json.h>
#include <gio/gio.h>

#include "utils/parser/fileparserfactory.h"
#include "utils/utils.h"

namespace {

const char* getTagPrompt = R"(
    # Role
    你是一个文档标签提取器，你的职责是针对给定的文档内容，按照要求抽取出标签，标签会用于用户在桌面操作系统重进行智能查找。

    # Context
    给定预置的文档标签分类：
    ···
    普通发票，增值税专用发票，餐饮发票，住宿发票，购物发票，交通发票，医疗发票，通讯发票，还款收据，承租收据，捐赠收据，购货收据，定金收据，押金收据，服务收据，付款收据，电子收据，医药费收据，报销归档，招待费用报销单，办公费用报销单，差旅费用报销单，通讯费用报销单，运输费用报销单，维修费用报销单，培训费用报销单，会议费用报销单，日常费用报表，差旅费用报表，员工培训费用报表，研发费用报表，设备维修与保养费用报表，项目费用报表，业绩报告，财务分析，财务指标，内部报告，预算报告，年度利润报表，年度审计报告，固定资产盘点表，库存资产盘点表，采购订单，维修记录，报废记录，固定期限合同，无固定期限合同，临时工合同，兼职合同，全职合同，入职协议，岗位说明书，保密协议，竞业协议，工作时间表，薪酬协议，绩效评估文件，培训协议，行政岗位入职登记表，研发岗位入职登记表，个人简历，月工资单，季度工资单，部门例会，项目进度，培训会议纪要，公关会议纪要，战略规划，工作总结，销售报告，项目研发报告，市场报告，人才管理报告，生产质量报告，信息采集，会议录制，红头文件，安全报告，政府工作报告，年假申请，调休假申请，事假申请，婚假申请，产假申请，陪产假申请，技术创新项目计划书，商业项目计划书，项目阶段性成果报告，项目中期进展报告，项目季度总结报告，项目进度概览报告，项目风险评估报告，技术风险评估报告，周期风险评估报告，自然灾害风险评估报告，社会风险评估报告，交通安全风险评估报告，信息风险评估报告，销售服务合同，委托服务合同，运输服务合同，租赁服务合同，咨询服务合同，房屋租赁协议，土地租赁协议，汽车租赁协议，仓库租赁协议，商业保密协议，雇佣保密协议，专利保密协议，机密信息保密协议，销售业绩报告，销售预测报告，销售费用报告，客户满意度与销售关系报告，市场竞争与销售策略报告，市场概况报告，竞争对手分析报告，产品市场调研报告，消费者调研报告，行业趋势分析报告，产品报价单，服务报价单，运输报价单，活动策划报价单，软件定制开发报价单，广告设计报价单，材料报价单，概要设计，详细设计，数据库设计文档，界面设计文档，应用图标设计，设计稿，平面设计，操作系统用户手册，软件用户手册，配置文件，测试用例，测试BUG，BUG修改记录，商业合同，劳动关系，税收，民事，刑事，产品反馈表，服务体验反馈表，客户忠诚度调查表，运营活动方案，市场调研问卷，政府与社会调查，组织管理调查问卷，教育评估问卷，个人年度总结，团队年度总结，公司年度总结
    ···

    # Flow
    按如下步骤工作：
    1. 预置标签匹配：结合<Context>给定的预置标签分类，输出给定文档内容最匹配的预置标签清单，最多5个；
    2. 自由标签提取：结合第一步的匹配结果，深入理解该类文档的关键特征，然后从用户检索文档的视角和习惯出发，输出最能能代表文档内容且符合用户记忆特点的标签，最多输出20个；

    # Rules
    - 标签通常是名词，避免输出其它时态的标签
    - 标签要简短且完整，少于15个汉字
    - 遇到多语言文档，需要把标签翻译成中文再输出
    - 预置标签只能按照原文输出，不可以更改和自由发挥
    - 自由标签应该优先抽取人物名称、地点、时间、公司/学校/部门名称、主题等便于人类记忆的内容
    - 只输出信心指数>0.7的标签

    # Output
    直接按如下JSON格式输出结果，不要解释：
    ···
    {
        "presetLabels": [{"label":"预置标签名", "confidence":"信心指数取值[0-1]"},{...}],
        "freeLabels":[{"label":"自由标签名", "confidence":"信心指数取值[0-1]"},{...}]
    }
    ···

    # Input
    给定的文档内容：
    ···
    ${docContent}
    ···

    # Initialize
    作为<Role>你需要仔细阅读<Context>和<Input>的内容，严格按照<Flow>指定的工作流程，遵循<Rules>定义的规则，按照<Output>的要求输出结果
)";

const char* getSummaryPrompt = R"(
    任务描述：为文档撰写简明摘要，不超过[字数限制]字，包含文档主要观点和关键信息。
    输入内容：‘[文档内容]’
    预期输出：文档摘要

    [文档内容]:
    ···
    ${docContent}
    ···
)";

}

std::string NlpHelper::chatSync(const std::string& message) const
{
    // 线程中启动一个子事件循环，通过子事件循环来处理回调
    // chat异步请求基于事件循序实现，如果主循环阻塞，则无法收到回调消息
    // 因此为了同步请求，需要在子循环中发送请求，阻塞主循环等待请求回复后退出子循环
    GMainLoop* subLoop = g_main_loop_new(nullptr, false);
    auto chatThread = std::async(std::launch::async, [this, subLoop, &message]() {
        std::string chatReply;

        // 在收到请求回调后退出子循环
        auto callback = [&chatReply, subLoop](const std::string& reply) {
            chatReply = reply;
            g_main_loop_quit(subLoop);
        };

        ky_aisdk::assistant::OsAssistant assistant;
        assistant.init();
        assistant.setChatAsyncCallback(callback);
        assistant.chatAsync(this->createChatRequestJson(message));
        g_main_loop_run(subLoop);

        return chatReply;
    });

    std::future_status status = chatThread.wait_for(std::chrono::seconds(60));
    if (status == std::future_status::timeout) {
        std::cerr << "Chat request timeout" << std::endl;
        g_main_loop_quit(subLoop);
        return std::string();
    }

    return chatThread.get();
}

std::vector<std::string> NlpHelper::getFileTagsSync(const std::string& filePath, const std::string& format) const
{
    std::string text = getDocContent(filePath, format);
    if (text.empty()) {
        std::cerr << "Get doc content failed." << std::endl;
        return std::vector<std::string>();
    }
    std::string prompt = getTagPrompt;
    addDocContentToPrompt(prompt, text);

    std::string reply = chatSync(prompt);
    if (reply.empty()) {
        std::cerr << "Get tags failed." << std::endl;
        return std::vector<std::string>();
    }

    std::vector<std::string> tags = extractTagsFromJson(reply);

    return tags;
}

std::string NlpHelper::getFileSummarySync(const std::string& filePath, const std::string& format) const
{
    std::string text = getDocContent(filePath, format);
    std::string prompt = getSummaryPrompt;
    addDocContentToPrompt(prompt, text);
    addMaxLengthToPrompt(prompt, 200);

    std::string reply = chatSync(prompt);
    if (reply.empty()) {
        std::cerr << "Get summary failed." << std::endl;
        return std::string();
    }

    std::string summary = extractSummaryFromJson(reply);
    summary = text::removeFirstLineWithPrefix(summary, "文档摘要：");

    return summary;
}

std::string NlpHelper::getDocContent(const std::string& filePath, const std::string& format) const
{
    if (!std::filesystem::exists(filePath)) {
        std::cerr << "File: " << filePath << " not exit, can't parser content." << std::endl;
        return std::string();
    }

    std::unique_ptr<Parser> parser = FileParserFactory::createParser(format);
    if (!parser) {
        std::cerr << "Can't create parser for " << format << std::endl;
        return std::string();
    }
    FileInfo fileInfo = parser->parse(filePath);

    // 将文件内容拼接成一个字符串
    std::ostringstream oss;
    for(const auto& str : fileInfo.contents) {
        oss << str;
    }
    std::string content = oss.str();

    return content;
}

std::string NlpHelper::getFileFormat(const std::string& filePath) const
{
    size_t dotIndex = filePath.find_last_of('.');
    if (dotIndex != std::string::npos) {
        return filePath.substr(dotIndex + 1);
    }

    return std::string(); // 如果未找到"."，返回空字符串表示无后缀
}

std::string NlpHelper::createChatRequestJson(const std::string& message) const
{
    Json::Value root;
    root["message_type"] = "chat";

    Json::Value contentArray;
    Json::Value content;
    content["type"] = "text";
    content["text"] = message;
    contentArray.append(content);
    root["content"] = contentArray;

    return root.toStyledString();
}

void NlpHelper::addDocContentToPrompt(std::string& prompt, const std::string& docContent) const
{
    const std::string replacedText = "${docContent}";
    prompt = prompt.replace(prompt.find(replacedText), replacedText.size(), docContent);
}

void NlpHelper::addMaxLengthToPrompt(std::string& prompt, int maxLength) const
{
    const std::string replacedText = "[字数限制]";
    prompt = prompt.replace(prompt.find(replacedText), replacedText.size(), std::to_string(maxLength));
}

std::vector<std::string> NlpHelper::extractTagsFromJson(const std::string& json) const
{
    std::vector<std::string> tags;

    Json::CharReaderBuilder builder;
    std::unique_ptr<Json::CharReader> reader(builder.newCharReader());
    Json::Value root;
    std::string errors;

    if (!reader->parse(json.c_str(), json.c_str() + json.size(), &root, &errors)) {
        std::cerr << "Failed to parse JSON: " << json << "\n" << errors << std::endl;
        return tags;
    }
    if (!root.isMember("content") || !root["content"].isArray() || root["content"].size() == 0) {
        std::cerr << "Failed to parse content key in json" << json << std::endl;
        return tags;
    }
    if (!root["content"][0].isMember("text") || !root["content"][0]["text"].isObject() || !root["content"][0]["text"].isMember("result")) {
        std::cerr << "Failed to parse text key in json" << json << std::endl;
        return tags;
    }

    // result 字段的值是一个字符串，格式是json，需要单独解析
    std::string resultJson = root["content"][0]["text"]["result"].asString();
    if (resultJson.empty()) {
        return tags;
    }
    Json::Value subRoot;
    if (!reader->parse(resultJson.c_str(), resultJson.c_str() + resultJson.size(), &subRoot, &errors)) {
        std::cerr << "Failed to parse result JSON: " << resultJson << "\n"
                  << "full JSON: " << json << "\n"
                  << errors << std::endl;
        return tags;
    }

    if (subRoot.isMember("freeLabels")) {
        for (const auto& label : subRoot["freeLabels"]) {
            std::string labelText = label["label"].asString();
            tags.push_back(labelText);
        }
    }
    if (subRoot.isMember("presetLabels")) {
        for (const auto& label : subRoot["presetLabels"]) {
            std::string labelText = label["label"].asString();
            tags.push_back(labelText);
        }
    }

    return tags;
}

std::string NlpHelper::extractSummaryFromJson(const std::string& json) const
{
    std::string summary;

    Json::CharReaderBuilder builder;
    std::unique_ptr<Json::CharReader> reader(builder.newCharReader());
    Json::Value root;
    std::string errors;

    bool parsingSuccessful = reader->parse(json.c_str(), json.c_str() + json.size(), &root, &errors);
    if (!parsingSuccessful) {
        std::cerr << "Failed to parse JSON: " << errors << std::endl;
        return summary;
    }

    summary = root["content"][0]["text"]["result"].asString();
    return summary;
}
